#!/usr/bin/python

# Copyright (C) 2009 Ori Mosenzon
#
# This program is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation; either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License along
# with this program; if not, write to the Free Software Foundation, Inc.,
# 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA.
# http://www.gnu.org/copyleft/gpl.html

import visual
import random

epsilone = 1e-7 

def find_index(obj,objs):
    i = 0 
    for o in objs:
        if o == obj:
            return i
        i +=1
    return None  


is_span_positive = False 

current_span = visual.curve()

span_frame = visual.frame() 

span_color = visual.color.yellow
span_fact = 1.2


def update_convex():
    global current_span
    current_span.visible = 0 

    poss = [orig] 
    sum = visual.vector(orig)

    vecs = map(lambda i:vectors[i].axis * span_fact, selected)

    poss = all_conv_comb(vecs)

    current_span = visual.convex(pos=poss, color = span_color)

def all_conv_comb(vecs):
    v0 = visual.vector(epsilone,epsilone,epsilone) # epsilone due to a bug in visual.convex

    if len(vecs) == 0:
        return [v0] 

    rec = all_conv_comb(vecs[:-1])

    v1 = vecs[-1] 

    ret = [] 

    if is_span_positive:
        vv = v0
    else:
        vv = -v1

    for s in rec:
        ret += [s+v1,s+vv]

    return ret 


def delete_current_grid(): 
    global span_frame
    for o in span_frame.objects:
        o.visible = 0 

def update_grid():

    delete_current_grid()

    global span_frame

    span_frame = visual.frame() 

    vecs = map(lambda i: visual.vector(vectors[i].axis)*span_fact,selected)
    for v in vecs:
        rest = vecs[0:len(vecs)] # deep copy  
        rest.remove(v)
        all_sums = all_sums_comb(rest)
        for aSum in all_sums:
            visual.curve(frame = span_frame, pos = [aSum,aSum+v], color = span_color)
            if not is_span_positive:
                visual.curve(frame = span_frame, pos = [aSum,aSum-v], color = span_color)

def all_sums_comb(vecs):
    vecs_segs = [] 
    res = 2 
    for v in vecs:
        norm = visual.mag(v)
        v1 = v*(res/norm)
        segs = [] 
        nn = int(norm/res)
        start = 0 
        if not is_span_positive:
            start = -nn
        for i in range(start,nn):
            segs += [v1*i]
        vecs_segs += [segs]
    cartecians = ordered_cartesian(vecs_segs)
    sums = [] 
    for c in cartecians:
        sum = visual.vector(orig)
        for v in c:
            sum += v 
        sums += [sum] 
    return sums     

def ordered_cartesian(list_of_lists):
    if not list_of_lists:
        return [[]] 
    aList = list_of_lists.pop()
    rec = ordered_cartesian(list_of_lists)
    ret = [] 
    for aCart in rec:
        for element in aList:
            ret += [aCart + [element]]
    return ret 


def is_selected(i):
    return vectors[i].color == select_color


## main ##

print """ 
Change view angle - Right drag 

Zoom in out       - Middle drag 

Select            - Left click on vector 

Move vctor        - Ledt drag a vector 

Grid span view    - 'g' on keyboard 

Convex span view  - 'c' on keyboard  

Full span         - Up arrow on keyboard 

Only positive span- Down arrow on keyboard

"""

span_vis = update_grid


orig = (0,0,0)
x_axis = (10,0,0)
y_axis = (0,10,0)
z_axis = (0,0,10)

rad = .1
select_color = visual.color.cyan
axes_color = visual.color.green

visual.cylinder(pos=orig, axis=x_axis, radius=rad, color = axes_color)
visual.cylinder(pos=orig, axis=y_axis, radius=rad, color = axes_color)
visual.cylinder(pos=orig, axis=z_axis, radius=rad, color = axes_color)

visual.label(pos=x_axis, text='X')
visual.label(pos=y_axis, text='Y')
visual.label(pos=z_axis, text='Z')

colors = [visual.color.red,visual.color.blue,visual.color.yellow,visual.color.orange,visual.color.white]
colors_num = len(colors)


vectors = []
cylinders = [] 

selected = [] # indexes of vectors

N = 4 
sw = .6
cyn_fact = .9


for i in range(N):
    ax = (random.randint(-10,10), random.randint(-10,10), random.randint(-10,10))

    v = visual.arrow(pos = orig, axis = ax, shaftwidth = sw, fixedwidth = 1, color = colors[i % colors_num])
    v_cyn = visual.cylinder(pos = orig, axis = visual.vector(ax)*cyn_fact, radius = sw/2, color = colors[i % colors_num]) 
    
    vectors += [v]
    cylinders += [v_cyn]

dragged = None 

while True:

    if dragged != None:
                ax = vectors[dragged].axis = visual.scene.mouse.pos
                cylinders[dragged].axis = visual.vector(ax)*cyn_fact

    if visual.scene.kb.keys:

        k = visual.scene.kb.getkey() 

        current_span.visible = 0

        if k == 'g':
            span_vis = update_grid
        elif k == 'c':
            delete_current_grid()
            span_vis = update_convex
        elif k=='down': 
            delete_current_grid()
            is_span_positive = True
        elif k=='up': 
            delete_current_grid()
            is_span_positive = False 

        span_vis()


    if not visual.scene.mouse.events:
        continue

    e = visual.scene.mouse.getevent()



    if e.drag:
        if dragged == None:
            dragged = find_index(e.pick,cylinders)
            if dragged == None:
                continue

    elif e.drop:
        if dragged != None:
            if is_selected(dragged):
                span_vis()
            dragged = None

    elif e.click:
        i = find_index(e.pick,cylinders)
        if i != None:
            if is_selected(i):
                vectors[i].color = colors[i % colors_num]
                selected.remove(i)
            else:
                selected += [i]
                vectors[i].color = select_color

            span_vis()
